package com.apop.demo.encrypt.util.sm;

import com.apop.demo.encrypt.util.HexUtils;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;


public class FzSM4 {
	private static final String ALGORITHM_NAME = "SM4";

	private static final int DEFAULT_KEY_SIZE = 64;

	private static final String ECB_MODE = "ECB";

	private static final String CBC_MODE = "CBC";

	private static Cipher generateEcbCipher(String algorithmName, int mode, byte[] key) throws Exception {
		Cipher cipher = Cipher.getInstance(algorithmName, "BC");
		Key sm4Key = new SecretKeySpec(key, "SM4");
		cipher.init(mode, sm4Key);
		return cipher;
	}

	private static Cipher generateCbcCipher(String algorithmName, int mode, byte[] key, byte[] bivSpec)
			throws Exception {
		Cipher cipher = Cipher.getInstance(algorithmName, "BC");
		Key sm4Key = new SecretKeySpec(key, "SM4");
		cipher.init(mode, sm4Key, generateIV(bivSpec));
		return cipher;
	}

	public static byte[] encryptEcb(byte[] hexKeyBytes, byte[] paramBytes, String algorithmName) throws Exception {
		Cipher cipher = generateEcbCipher(algorithmName, 1, hexKeyBytes);
		return cipher.doFinal(paramBytes);
	}

	public static byte[] encryptCbc(byte[] hexKeyBytes, byte[] paramBytes, String algorithmName, byte[] bivSpec)
			throws Exception {
		Cipher cipher = generateCbcCipher(algorithmName, 1, hexKeyBytes, bivSpec);
		return cipher.doFinal(paramBytes);
	}

	public static byte[] encrypt(byte[] hexKeyBytes, byte[] paramBytes, String algorithmName, byte[] bivSpec)
			throws Exception {
		byte[] result;
		String[] algorithmList = algorithmName.split("/");
		String mode = algorithmList[1];
		if ("ECB".equals(mode.toUpperCase())) {
			result = encryptEcb(hexKeyBytes, paramBytes, algorithmName);
		} else {
			if (!"CBC".equals(mode.toUpperCase())) {
				throw new IllegalArgumentException("algorithmName parameter is illegal,can parse ecb or cbc! ");
			}
			result = encryptCbc(hexKeyBytes, paramBytes, algorithmName, bivSpec);
		}
		return result;
	}

	public static byte[] decryptEcb(byte[] hexKeyBytes, byte[] cipherBytes, String algorithmName) throws Exception {
		Cipher cipher = generateEcbCipher(algorithmName, 2, hexKeyBytes);
		return cipher.doFinal(cipherBytes);
	}

	public static byte[] decryptCbc(byte[] hexKeyBytes, byte[] cipherBytes, String algorithmName, byte[] bivSpec)
			throws Exception {
		Cipher cipher = generateCbcCipher(algorithmName, 2, hexKeyBytes, bivSpec);
		return cipher.doFinal(cipherBytes);
	}

	public static byte[] decrypt(byte[] hexKeyBytes, byte[] cipherBytes, String algorithmName, byte[] bivSpec)
			throws Exception {
		byte[] result;
		String[] algorithmList = algorithmName.split("/");
		String mode = algorithmList[1];
		if ("ECB".equals(mode.toUpperCase())) {
			result = decryptEcb(hexKeyBytes, cipherBytes, algorithmName);
		} else {
			if (!"CBC".equals(mode.toUpperCase())) {
				throw new IllegalArgumentException("algorithmName parameter is illegal,can parse ecb or cbc! ");
			}
			result = decryptCbc(hexKeyBytes, cipherBytes, algorithmName, bivSpec);
		}
		return result;
	}

	public static String generateKey() throws Exception {
		KeyGenerator kg = null;
		kg = KeyGenerator.getInstance("SM4", "BC");
		kg.init(64, new SecureRandom());
		byte[] keyBytes = kg.generateKey().getEncoded();
		return HexUtils.byteToHex(keyBytes).toUpperCase();
	}

	public static AlgorithmParameters generateIV(byte[] iv) throws Exception {
		if (iv == null) {
			throw new IllegalArgumentException("IV parameter can't be null! ");
		}
		AlgorithmParameters params = AlgorithmParameters.getInstance("SM4");
		params.init(new IvParameterSpec(iv));
		return params;
	}

	static {
		if (Security.getProvider("BC") == null) {
			Security.addProvider((Provider) new BouncyCastleProvider());
		}
	}
}
